from transformers import BertTokenizer,BertModel
import pandas as pd
import json
import pickle

def text_processing(text,bert_path='../bert-base-chinese'):
    tokenizer = BertTokenizer.from_pretrained(bert_path)
    bert = BertModel.from_pretrained(bert_path)
    text_tokenizer = tokenizer(text, padding = True, truncation = True, max_length = 512,
                               return_tensors = 'pt')
    output = bert(input_ids = text_tokenizer['input_ids'], attention_mask = text_tokenizer['attention_mask'],
                                 token_type_ids = text_tokenizer['token_type_ids'],
                                 return_dict=False)
    #torch.Size([1, 498, 768])->pooled_output[0]
    #torch.Size([1, 768])->pooled_output[-1]
    return output[-1][0].detach().numpy() #均值编码,转化数组

if __name__=='__main__':
    data = pd.read_csv('../data/lyrics.csv')
    text_list = []
    for index, text in enumerate(data['歌词']):
        dict_ = {}
        dict_['id'] = index
        dict_['text'] = text
        dict_['text_embedding'] = text_processing(text)
        text_list.append(dict_)
    # 保存列表数据到文件
    with open('data.pkl', 'wb') as f:
        pickle.dump(text_list, f)
